import sys

sys.setrecursionlimit(10000)


class GomokuAI:
    def __init__(self, n=15):
        self.n = n
        self.board = [[-1 for _ in range(n)] for _ in range(n)]
        self.inf = float('inf')
        self.directions = [(0, 1), (1, 0), (1, 1), (1, -1)]  # 右、下、右下、左下

    def check_win(self, player, k=5):
        """
        检查 player 是否在 board 中形成连续 k 子
        :param player: 玩家标识 (0 表示电脑, 1 表示玩家)
        :param k: 连子数，默认为5
        :return: True 如果获胜，否则 False
        """
        rows, cols = self.n, self.n

        for i in range(rows):
            for j in range(cols):
                if self.board[i][j] != player:
                    continue
                # 从 (i,j) 出发检测 4 个方向
                for dx, dy in self.directions:
                    count = 1  # 当前位置已经算一个
                    x, y = i + dx, j + dy
                    # 向一个方向延伸
                    while 0 <= x < rows and 0 <= y < cols and self.board[x][y] == player:
                        count += 1
                        if count == k:
                            return True
                        x += dx
                        y += dy
        return False

    def count_consecutive(self, player, row, col, dx, dy):
        """计算在给定方向上连续的棋子数"""
        count = 0
        x, y = row, col
        while 0 <= x < self.n and 0 <= y < self.n and self.board[x][y] == player:
            count += 1
            x += dx
            y += dy
        return count

    def evaluate_line(self, player, row, col, dx, dy):
        """评估一条线的得分"""
        # 检查以当前位置为中心的连线
        count = 1  # 当前位置
        blocked = 0  # 被堵住的端点数

        # 向一个方向检查
        x, y = row + dx, col + dy
        while 0 <= x < self.n and 0 <= y < self.n and self.board[x][y] == player:
            count += 1
            x += dx
            y += dy
        # 检查是否被堵住
        if not (0 <= x < self.n and 0 <= y < self.n) or self.board[x][y] != -1:
            blocked += 1

        # 向相反方向检查
        x, y = row - dx, col - dy
        while 0 <= x < self.n and 0 <= y < self.n and self.board[x][y] == player:
            count += 1
            x -= dx
            y -= dy
        # 检查是否被堵住
        if not (0 <= x < self.n and 0 <= y < self.n) or self.board[x][y] != -1:
            blocked += 1

        # 根据连子数和堵住情况评分
        if count >= 5:
            return 100000
        elif count == 4 and blocked == 0:
            return 10000
        elif count == 4 and blocked == 1:
            return 1000
        elif count == 3 and blocked == 0:
            return 100
        elif count == 3 and blocked == 1:
            return 10
        elif count == 2 and blocked == 0:
            return 5
        else:
            return 1

    def evaluate_position(self, player, row, col):
        """评估在某个位置落子的价值"""
        if self.board[row][col] != -1:
            return -self.inf  # 无效位置

        score = 0
        # 在四个方向上评估
        for dx, dy in self.directions:
            score += self.evaluate_line(player, row, col, dx, dy)
        return score

    def evaluate_board(self):
        """评估整个棋盘的局势"""
        score = 0
        for i in range(self.n):
            for j in range(self.n):
                if self.board[i][j] == -1:  # 空位
                    # 评估电脑在此落子的得分
                    ai_score = self.evaluate_position(0, i, j)
                    # 评估玩家在此落子的得分
                    player_score = self.evaluate_position(1, i, j)
                    # 取较大值作为该位置的价值
                    score += ai_score - player_score * 0.9  # 稍微优先防守
        return score

    def is_board_full(self):
        """检查棋盘是否已满"""
        for i in range(self.n):
            for j in range(self.n):
                if self.board[i][j] == -1:
                    return False
        return True

    def minimax(self, player, depth, alpha, beta):
        """
        带Alpha-Beta剪枝的极小化极大算法
        :param player: 当前玩家 (0 表示电脑, 1 表示玩家)
        :param depth: 搜索深度
        :param alpha: alpha值
        :param beta: beta值
        :return: (最佳评分, 最佳位置)
        """
        # 终止条件：有玩家获胜或达到搜索深度或棋盘已满
        if self.check_win(0):
            return 1000000, None
        if self.check_win(1):
            return -1000000, None
        if depth == 0 or self.is_board_full():
            return self.evaluate_board(), None

        best_move = None
        if player == 0:  # 电脑回合，最大化
            max_val = -self.inf
            # 优先考虑中心附近的位置
            moves = []
            center = self.n // 2
            for i in range(self.n):
                for j in range(self.n):
                    if self.board[i][j] == -1:
                        # 根据到中心的距离排序，优先考虑中心附近
                        distance = abs(i - center) + abs(j - center)
                        moves.append((distance, i, j))
            moves.sort()  # 距离越小越优先

            for _, i, j in moves:
                self.board[i][j] = 0
                val, _ = self.minimax(1, depth - 1, alpha, beta)
                self.board[i][j] = -1

                if val > max_val:
                    max_val = val
                    best_move = (i, j)
                alpha = max(alpha, val)
                if beta <= alpha:
                    break  # Alpha-Beta剪枝
            return max_val, best_move
        else:  # 玩家回合，最小化
            min_val = self.inf
            moves = []
            center = self.n // 2
            for i in range(self.n):
                for j in range(self.n):
                    if self.board[i][j] == -1:
                        distance = abs(i - center) + abs(j - center)
                        moves.append((distance, i, j))
            moves.sort()

            for _, i, j in moves:
                self.board[i][j] = 1
                val, _ = self.minimax(0, depth - 1, alpha, beta)
                self.board[i][j] = -1

                if val < min_val:
                    min_val = val
                    best_move = (i, j)
                beta = min(beta, val)
                if beta <= alpha:
                    break  # Alpha-Beta剪枝
            return min_val, best_move

    def get_best_move(self, depth=4):
        """获取电脑的最佳落子位置"""
        _, move = self.minimax(0, depth, -self.inf, self.inf)
        return move

    def print_board(self):
        """打印棋盘"""
        print("  ", end="")
        for j in range(self.n):
            print(f"{j:2}", end=" ")
        print()
        for i in range(self.n):
            print(f"{i:2}", end=" ")
            for j in range(self.n):
                if self.board[i][j] == -1:
                    print(" .", end=" ")
                elif self.board[i][j] == 0:
                    print(" O", end=" ")  # 电脑棋子
                else:
                    print(" X", end=" ")  # 玩家棋子
            print()

    def play_game(self):
        """开始游戏"""
        print(f"欢迎来到 {self.n}x{self.n} 五子棋游戏！")
        print("你是 X，电脑是 O")
        print("输入坐标格式：行 列 (从0开始)")

        while True:
            self.print_board()

            # 玩家回合
            while True:
                try:
                    print("请输入你的落子位置 (行 列): ")
                    x, y = map(int, input().split())
                    if 0 <= x < self.n and 0 <= y < self.n and self.board[x][y] == -1:
                        self.board[x][y] = 1
                        break
                    else:
                        print("无效位置，请重新输入！")
                except:
                    print("输入格式错误，请重新输入！")

            # 检查玩家是否获胜
            if self.check_win(1):
                self.print_board()
                print("恭喜你获胜！")
                break

            # 检查是否平局
            if self.is_board_full():
                self.print_board()
                print("平局！")
                break

            # 电脑回合
            print("电脑思考中...")
            move = self.get_best_move()
            if move:
                self.board[move[0]][move[1]] = 0
                print(f"电脑落子位置: {move[0]} {move[1]}")

            # 检查电脑是否获胜
            if self.check_win(0):
                self.print_board()
                print("电脑获胜！")
                break

            # 检查是否平局
            if self.is_board_full():
                self.print_board()
                print("平局！")
                break


# 使用示例
if __name__ == "__main__":
    # 可以指定棋盘大小，默认为15
    game = GomokuAI(5)  # 使用11x11棋盘，适合测试
    game.play_game()
